#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从JSON文件导入分类数据到数据库
"""

import sqlite3
import json
import os
import sys

def get_db_path():
    """获取数据库路径"""
    current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    db_path = os.path.join(current_dir, "..", "flux_prompt_helper", "promptManager-server", "instance", "prompt_manager.db")
    return db_path

def create_categories_table(cursor):
    """创建分类表"""
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS categories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            category_key TEXT NOT NULL UNIQUE,
            category_name TEXT NOT NULL,
            description TEXT DEFAULT '',
            parent_id INTEGER DEFAULT NULL,
            level INTEGER DEFAULT 0,
            data TEXT NOT NULL,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (parent_id) REFERENCES categories (id)
        )
    """)
    
    # 创建索引
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_categories_key ON categories (category_key)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_categories_parent ON categories (parent_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_categories_level ON categories (level)")

def import_from_json():
    """从JSON文件导入分类数据"""
    # 读取JSON文件
    current_dir = os.path.dirname(os.path.abspath(__file__))
    json_path = os.path.join(current_dir, "..", "enhanced_categories.json")
    
    if not os.path.exists(json_path):
        print(f"JSON文件不存在: {json_path}")
        return False
    
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 获取数据库路径
    db_path = get_db_path()
    if not os.path.exists(os.path.dirname(db_path)):
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
    
    # 连接数据库
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 创建表
        create_categories_table(cursor)
        
        # 清空现有数据
        cursor.execute("DELETE FROM categories")
        print("清空现有分类数据")
        
        # 处理分类数据
        categories = data.get('categories', {})
        
        print(f"开始导入 {len(categories)} 个主分类")
        
        for category_key, category_data in categories.items():
            # 插入主分类
            cursor.execute("""
                INSERT INTO categories 
                (category_key, category_name, description, parent_id, level, data, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
            """, (
                category_key,
                category_data['name'],
                category_data['description'],
                None,
                0,
                json.dumps(category_data, ensure_ascii=False)
            ))
            
            # 获取主分类ID
            cursor.execute("SELECT id FROM categories WHERE category_key = ?", (category_key,))
            main_category_id = cursor.fetchone()[0]
            
            print(f"导入主分类: {category_data['name']}")
            
            # 处理子分类
            subcategories = category_data.get('subcategories', {})
            for sub_key, sub_data in subcategories.items():
                cursor.execute("""
                    INSERT INTO categories 
                    (category_key, category_name, description, parent_id, level, data, created_at, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
                """, (
                    f"{category_key}_{sub_key}",
                    sub_data['name'],
                    sub_data['description'],
                    main_category_id,
                    1,
                    json.dumps(sub_data, ensure_ascii=False)
                ))
                
                print(f"  导入子分类: {sub_data['name']}")
        
        # 提交事务
        conn.commit()
        
        # 统计导入结果
        cursor.execute("SELECT COUNT(*) FROM categories")
        total_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM categories WHERE level = 0")
        main_count = cursor.fetchone()[0]
        
        cursor.execute("SELECT COUNT(*) FROM categories WHERE level > 0")
        sub_count = cursor.fetchone()[0]
        
        print("导入完成!")
        print(f"总计: {total_count} 个分类")
        print(f"主分类: {main_count} 个")
        print(f"子分类: {sub_count} 个")
        
        return True
        
    except Exception as e:
        print(f"导入失败: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

def main():
    """主函数"""
    print("开始从JSON文件导入分类数据")
    print("=" * 50)
    
    success = import_from_json()
    
    if success:
        print("分类数据导入成功!")
    else:
        print("分类数据导入失败!")

if __name__ == "__main__":
    main()
